import Model.AutoEncoding as MyModel # 模型
import Loader as MyLoader # 数据加载器
import Config # 训练参数
import PSNR # 评价图像相似程度，越大越相似

import torch
from torch import nn,optim
from torch.utils.tensorboard import SummaryWriter
# tensorboard --logdir={Config.SUMMARY_WRITER_PATH} --port 8123
writer = SummaryWriter(Config.SUMMARY_WRITER_PATH)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


model = MyModel.DenoiseAutoEncoder().to(device)
train_data_loader,val_data_loader = MyLoader.loadTrainValData()

criterion = nn.MSELoss().to(device)
optimizer = optim.Adam(model.parameters(), lr=Config.LR)

min_val_loss = 0 # 现有最小验证集损失
# 实际训练
for epoch in range(Config.MAX_EPOCH):

    print(epoch+1)
    model = model.train()
    batch_num = 0
    total_loss = 0
    for batch_data in train_data_loader:
        optimizer.zero_grad()

        X = batch_data['X']
        y = batch_data['y']

        outputs = model(X)
        loss = criterion(outputs, y)

        total_loss += loss.item() * X.size(0)
        batch_num += X.size(0)

        loss.backward()
        optimizer.step()
        
        
    writer.add_scalar(f"Ave Train MSELoss",total_loss/batch_num,epoch)

    if (epoch+1) % 5 == 1:
        model = model.eval()
        batch_num = 0
        total_loss = 0
        for batch_data in val_data_loader:
            X = batch_data['X']
            y = batch_data['y']

            outputs = model(X)
            loss = criterion(outputs, y)

            total_loss += loss.item() * X.size(0)
            batch_num += X.size(0)

            # torch.save({'model': model.state_dict()}, f'./model/model_epoch_{epoch}.pth')
        writer.add_scalar(f"Ave Validate MSELoss",total_loss/batch_num,epoch)
        print(f"epoch {epoch+1} 的Ave Validate MSELoss为{total_loss/batch_num}")
        # 选择Loss最小进行存储
        if (total_loss/batch_num < min_val_loss):
            min_val_loss = total_loss/batch_num
            torch.save({'model': model.state_dict()}, Config.BEST_MODEL_PATH)
        
